package com.iddactylogram.client;

import com.iddactylogram.common.domain.IdConstant;
import com.iddactylogram.common.domain.NetworkAddress;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.api.CuratorWatcher;
import org.apache.curator.framework.imps.CuratorFrameworkState;
import org.apache.curator.retry.RetryNTimes;
import org.apache.zookeeper.WatchedEvent;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.function.Consumer;
import java.util.logging.Level;
import java.util.logging.Logger;

public class ZkLoadBalancing implements ILoadBalancing {
    private final static Logger LOG = Logger.getLogger(ZkLoadBalancing.class.getName());
    private final CuratorFramework curatorFramework;
    private final ConcurrentMap<String, Long> serverInfoCache = new ConcurrentHashMap<>();
    private Consumer<Object> consumer = null;

    public ZkLoadBalancing(CuratorFramework curatorFramework) {
        CuratorFrameworkState state = curatorFramework.getState();
        if (state != CuratorFrameworkState.STARTED) {
            curatorFramework.start();
        }
        this.curatorFramework = curatorFramework.usingNamespace(IdConstant.DEFAULT_ZK_NAMESPACE);
    }

    public ZkLoadBalancing(String zkAddress) {
        CuratorFramework temp = CuratorFrameworkFactory.newClient(zkAddress, new RetryNTimes(3, 1000));
        temp.start();
        curatorFramework = temp.usingNamespace(IdConstant.DEFAULT_ZK_NAMESPACE);
        addWatch();
    }

    @Override
    public void loadServerInfo(Consumer<Object> balancingAfter) {
        try {
            this.consumer = balancingAfter;
            serverInfoCache.clear();
            List<String> serverInfos = curatorFramework.getChildren().forPath("/");
            serverInfos.forEach(server -> {
                serverInfoCache.put(server, System.currentTimeMillis());
            });
        } catch (Exception e) {
            LOG.log(Level.INFO, "loadServerInfo fail:%s", e);
            throw new IllegalStateException("loadServerInfo fail", e);
        }
    }

    @Override
    public NetworkAddress chooseServerAddress(String appName) {
        int size = serverInfoCache.size();
        int choose = Math.abs(appName.hashCode() % size);
        Set<Map.Entry<String, Long>> entries = serverInfoCache.entrySet();
        int index = 0;
        for (Map.Entry<String, Long> entry : entries) {
            if (choose == (index++)) {
                String serverAddress = entry.getKey();
                String[] ss = serverAddress.split(":");
                return new NetworkAddress(ss[0], Integer.parseInt(ss[1]));
            }
        }
        String msg = String.format("没有选择到合适的server，appName:%s, appNameHashCode:%s, serverInfoCache:%s", appName, choose, serverInfoCache);
        throw new IllegalStateException(msg);
    }

    private void addWatch() {
        try {
            this.curatorFramework.getChildren().usingWatcher(new IdCuratorWatcher()).forPath("/");
        } catch (Exception e) {
            LOG.log(Level.INFO, "addWatch fail :%s", e);
        }
    }

    private class IdCuratorWatcher implements CuratorWatcher {
        @Override
        public void process(WatchedEvent event) {
            System.out.println("监听ZK的事件，进行从新获取服务器信息处理");
            //重新再次注册监听
            addWatch();
            //发送通知，马上断开从新链接
            if (consumer != null) {
                consumer.accept("重新加载");
            }
        }
    }
}
